package ws2

import (
	"fmt"
	"net/http"
	"crypto/tls"
	"github.com/gorilla/websocket"
	"sync"
	// "time"
	"strings"
	"runtime/debug"

	"gitee.com/wint/tge"
)

const (
	CONN_BUF_FULL_CLOSE = 0		// 发送缓冲区满关闭连接
	CONN_BUF_FULL_DISCARD = 1	// 发送缓冲区满丢弃发送数据
)

// 在连接的发送缓冲区满了之后的行为
var ConnBuffFullAction = CONN_BUF_FULL_CLOSE

type WebSocketServer struct {
	Conns		map[string]*websocket.Conn
	CreateAgent	func(tge.Conn) tge.Agent

	// https 相关配置
	HttpsAddr 	string
	CertFile 	string
	KeyFile 	string

	Handler 	*http.ServeMux 	// 响应接口

	server 		*http.Server
	serverHttps *http.Server

	IsStarted	bool
	mu 			sync.RWMutex
	wg 			sync.WaitGroup
}

func NewWebSocketServer() *WebSocketServer {
	s := &WebSocketServer{}
	s.Conns		= make(map[string]*websocket.Conn)
	s.Handler 	= http.NewServeMux()

	upgrader := websocket.Upgrader {
		// 允许跨域
		CheckOrigin:func(r *http.Request) bool{
			return true
		},
	}
	s.Handler.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
		if wsConn, err := upgrader.Upgrade(w, r, nil); err == nil{
			s.wsConnHandler(wsConn)
		}
	})

	s.server = &http.Server{}
	s.IsStarted = false
	return s
}

func (s *WebSocketServer) SetCreateAgent (fn func(tge.Conn) tge.Agent) {
	s.CreateAgent = fn
}

func (s *WebSocketServer) wsConnHandler(conn *websocket.Conn) {
	defer func() {
		if err := recover(); err != nil {
			tge.Logger.Println("wsConnHandler", err, string(debug.Stack()))
		}
	}()

	key := conn.RemoteAddr().String()
	count := 0

	s.mu.Lock()
	s.Conns[key] = conn
	count = len(s.Conns)
	s.mu.Unlock()
	
	tge.Logger.Println("Accept ws from", key, count)

	defer conn.Close()


	s.wg.Add(1)
	defer s.wg.Done()

	agent := s.CreateAgent(NewWebSocketConn(conn))
	defer agent.Close()
	agent.Run()

	s.mu.Lock()
	delete(s.Conns, key)
	count = len(s.Conns)
	s.mu.Unlock()

	tge.Logger.Println("close ws from", key, count)
}

func (s *WebSocketServer) StartServer(addr string) {
	if "" != s.HttpsAddr && "" != s.CertFile && "" != s.KeyFile {
		go s.StartServerTLS(s.HttpsAddr, s.CertFile, s.KeyFile)
	}

	if !strings.Contains(addr, ":") {
		addr = ":" + addr
	}
	tge.Logger.Println("Listen Websocket on http ", addr)
	s.IsStarted = true
	s.server.Addr = addr
	s.server.Handler = s.Handler
	if err := s.server.ListenAndServe(); nil != err {
		fmt.Println("[WebSocketServer] ListenAndServeTLS error", err)
	}
}

func (s *WebSocketServer) StartServerTLS(addr string, certFile string, keyFile string) {
	if !strings.Contains(addr, ":") {
		addr = ":" + addr
	}
	tge.Logger.Println("Listen Websocket on https ", addr)
	cfg := &tls.Config{
        // MinVersion: tls.VersionTLS12,
    }
	s.serverHttps = &http.Server{
		TLSConfig: cfg,
	}
	s.serverHttps.Addr = addr
	s.serverHttps.Handler = s.Handler
	if err := s.serverHttps.ListenAndServeTLS(certFile, keyFile); nil != err {
		fmt.Println("[WebSocketServer] ListenAndServeTLS error", err)
	}
}

func (s *WebSocketServer) Close() {
	if !s.IsStarted {
		tge.Logger.Println("WebSocketServer isn't started")
		return
	}

	s.IsStarted = false

	s.server.Close()

	if nil != s.serverHttps {
		s.serverHttps.Close()
	}

	// t := time.Now().Add(100 * time.Millisecond)
	s.mu.Lock()
	for _, v := range s.Conns {
		// v.SetDeadline(t)
		v.Close()
	}
	s.mu.Unlock()

	tge.Logger.Println("WebSocketServer waiting for threads end")
	s.wg.Wait()

	tge.Logger.Println("WebSocketServer is closed")
}
